from .ConvLSTM import ConvLSTM
from .Encode import Encode
from .Res_MCNN import Res_MCNN
from .GCN import GCN
import torch
import torch.nn as nn
class RF_STED(nn.Module):
    def __init__(self):
        super(RF_STED,self).__init__()
        self.covlstm = ConvLSTM(input_dim = 1, 
                                hidden_dim = 32, 
                                kernel_size = (3,3),
                                num_layers = 1,
                                batch_first=True, 
                                bias=True, 
                                return_all_layers=False)
        self.gcn1 = GCN(in_features = 6,out_features = 1)   #输入3个时间步，输出3个特征
        self.gcn2 = GCN(in_features = 6,out_features = 1)
        self.gcn3 = GCN(in_features = 6,out_features = 1)
        self.gcn4 = GCN(in_features = 6,out_features = 1)
        self.encode = Encode()
        self.Res_CNN1 = Res_MCNN(num_convolutions=3, kernel_size=31, scale=2)
        self.Res_CNN2 = Res_MCNN(num_convolutions=3, kernel_size=33, scale=2)
        self.Res_CNN3 = Res_MCNN(num_convolutions=3, kernel_size=35, scale=2)   #num_convolutions越小越好  kernel_size居中  scale 2 差不多
        self.relu = nn.ReLU(inplace=True) #relu激活
    def forward(self,x,adj1,adj2,adj3,adj4):  #前向传播   x:[6,1,6,69,69]  adj:[4761,4761]
        #时间
        taxi_three_time_x = x[:,0,0:3,:,:]   #[6, 3, 69, 69]
        taxi_three_day_x = x[:,0,3:6,:,:]    #[6, 3, 69, 69]
        taxi_three_time_x = torch.unsqueeze(taxi_three_time_x,2) #[6, 3, 1,69, 69]
        taxi_three_day_x = torch.unsqueeze(taxi_three_day_x,2)   #[6, 3, 1,69, 69]
        taxi_three_time_x = self.covlstm(taxi_three_time_x)[1][0][0]  #[6,32,69,69]
        taxi_three_day_x = self.covlstm(taxi_three_day_x)[1][0][0]   #[6,32,69,69]
        taxi_three_time_x = taxi_three_time_x.permute(0,2,3,1)  #[6,69,69,32]
        taxi_three_day_x = taxi_three_day_x.permute(0,2,3,1)   #[6,69,69,32]
        #空间
        sp =  x.squeeze(1)#(6,6,69,69)
        sp = sp.view((sp.shape[0],sp.shape[1],sp.shape[2]*sp.shape[3]))#(6,6,4761)
        sp = sp.permute(0,2,1)                           #torch.Size([6, 4761, 6])
        sp1 = self.gcn1(x = sp,adj = adj1).view(sp.shape[0],69,69,-1)     
        sp2 = self.gcn1(x = sp,adj = adj2).view(sp.shape[0],69,69,-1) 
        sp3 = self.gcn1(x = sp,adj = adj3).view(sp.shape[0],69,69,-1) 
        sp4 = self.gcn1(x = sp,adj = adj4).view(sp.shape[0],69,69,-1)    #[6，69,69，1]
        # 编码
        time = torch.cat((taxi_three_time_x.unsqueeze(-1), taxi_three_day_x.unsqueeze(-1)), dim=-1)         #[6,69,69,32,2]
        time = time.view(time.shape[0], time.shape[1], time.shape[2] , time.shape[3] * time.shape[4])  #[6,69,69,64]
        time = time.sum(dim=3, keepdim=True)  #[6,69,69,1]
        space = torch.cat((sp1.unsqueeze(-1), sp2.unsqueeze(-1), sp3.unsqueeze(-1), sp4.unsqueeze(-1)), dim=-1) #[6,69,69,1,4]
        space = sp2
        encoder = time
        # encoder = self.encode(time,space)  #[6,69,69,1]
        # 解码
        encoder = encoder.permute(0,3,1,2)   #[6,1,69,69]
        encoder = self.Res_CNN1(encoder) #[6,1,69,69]
        encoder = self.Res_CNN2(encoder) #[6,1,69,69]
        encoder = self.Res_CNN3(encoder) #[6,1,69,69]
        out = self.relu(encoder)
        return out
